use rand::{thread_rng, Rng};
use std::mem;

/// 计算 (x * y + c) % m
fn mul_mod(x: u128, y: u128, c: u128, m: u128) -> u128 {
    if let Some(z) = x.checked_mul(y) {
        return (z + c) % m;
    }

    let mut ret = c;
    let mut x = x % m;
    let mut y = y % m;
    if x < y {
        mem::swap(&mut x, &mut y);
    }
    while y > 0 {
        if y & 1 == 1 {
            ret = (ret + x) % m;
        }
        x = (x << 1) % m;
        y >>= 1;
    }
    ret
}

/// 计算 (x ^ y) % m
fn pow_mod(mut x: u128, mut y: u128, m: u128) -> u128 {
    let mut ret = 1;
    while y > 0 {
        if y & 1 == 1 {
            ret = mul_mod(ret, x, 0, m);
        }
        x = mul_mod(x, x, 0, m);
        y >>= 1;
    }
    ret
}

/// 计算 x 和 y 的最大公约数
fn gcd(mut x: u128, mut y: u128) -> u128 {
    while y > 0 {
        let t = x % y;
        x = y;
        y = t;
    }
    x
}

/// 判断 x 是否为素数
fn is_prime(x: u128) -> bool {
    // 对小质数和常见合数的快速检查路径
    if x <= 5 {
        return x > 1 && x != 4;
    }
    if x % 2 == 0 || x % 3 == 0 || x % 5 == 0 {
        return false;
    }

    // Miller-Rabin 素性测试
    let k = (x - 1).trailing_zeros();
    let t = (x - 1) >> k;
    let mr_test = |x: u128, probe: u128| -> bool {
        let mut cur = pow_mod(probe, t, x);
        for _ in 0..k {
            let nxt = mul_mod(cur, cur, 0, x);
            if nxt == 1 && cur != 1 && cur != x - 1 {
                return false;
            }
            cur = nxt;
        }
        cur == 1
    };

    // https://miller-rabin.appspot.com/
    match x {
        x if x < 341531 => mr_test(x, 9345883071009581737),
        x if x < 350269456337 => [
            4230279247111683200,
            14694767155120705706,
            16641139526367750375,
        ]
        .into_iter()
        .all(|probe| mr_test(x, probe)),
        x if x < 318665857834031151167461 => {
            let probes = [
                2,
                4130806001517,
                149795463772692060,
                186635894390467037,
                396730417934771580,
                801385935568343050,
            ];
            probes.into_iter().all(|probe| mr_test(x, probe))
        }
        _ => panic!("unsupported range"),
    }
}

/// 基于 Brent 判环的 Pollard-Rho
fn pollard_rho(x: u128) -> u128 {
    let mut s = 0;
    let mut t = 0;
    let c = thread_rng().gen_range(1..x);

    // 每次倍增 steps，以减少 gcd 的调用次数
    let mut steps = 1;
    loop {
        let mut val = 1;
        for step in 1..=steps {
            t = mul_mod(t, t, c, x);
            val = mul_mod(val, t.abs_diff(s), 0, x);
            // 倍增优化
            if step % 127 == 0 {
                let d = gcd(val, x);
                if d > 1 {
                    return d;
                }
            }
        }
        let d = gcd(val, x);
        if d > 1 {
            return d;
        }
        steps <<= 1;
        s = t;
    }
}

/// 计算 max(n 的最大质因子, c)
pub fn find_prime_factor_lt(mut n: u128, mut c: u128) -> u128 {
    if n == 1 || n < c {
        return c;
    }
    if is_prime(n) {
        return n;
    }
    let mut p = n;
    while p >= n {
        p = pollard_rho(n);
    }
    while n % p == 0 {
        n /= p;
    }

    // 启发式方法，尽量更新 c 为较大的质因子
    c = find_prime_factor_lt(n.max(p), c);
    c = find_prime_factor_lt(n.min(p), c);
    c
}

pub fn find_max_prime_factor(n: u128) -> u128 {
    find_prime_factor_lt(n, 0)
}
